import argparse
import heapq

# Helper function to load sample keys from a TSV file
def load_sample_keys(file_path):
    sample_keys = set()
    with open(file_path, 'r') as file:
        for line in file:
            columns = line.strip().split('\t')
            sample_keys.add('\t'.join(columns[:3]))
    return sample_keys

# Helper function to process predictions and filter top samples
def process_predictions(predictions_file, pos_keys, neg_keys, top_n):
    pos_samples = []
    neg_samples = []

    with open(predictions_file, 'r') as file:
        for line in file:
            columns = line.strip().split('\t')
            base_key = '\t'.join(columns[:3])
            full_key = '\t'.join(columns[:3]) + '\t' + columns[4]  # Include the fifth column for the full key
            pred_0 = float(columns[6])
            pred_1 = float(columns[7])

            if base_key in pos_keys:
                heapq.heappush(pos_samples, (pred_1, full_key))
                if len(pos_samples) > top_n:
                    heapq.heappop(pos_samples)
            elif base_key in neg_keys:
                heapq.heappush(neg_samples, (pred_0, full_key))
                if len(neg_samples) > top_n:
                    heapq.heappop(neg_samples)

    return pos_samples, neg_samples

# Main function
def main(pos_file, neg_file, predictions_file, pos_output, neg_output, top_n):
    pos_keys = load_sample_keys(pos_file)
    neg_keys = load_sample_keys(neg_file)

    pos_samples, neg_samples = process_predictions(predictions_file, pos_keys, neg_keys, top_n)

    with open(pos_output, 'w') as pos_out:
        for _, key in sorted(pos_samples, reverse=True):
            pos_out.write(key + '\n')

    with open(neg_output, 'w') as neg_out:
        for _, key in sorted(neg_samples, reverse=True):
            neg_out.write(key + '\n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Filter top predictions for positive and negative samples.")
    parser.add_argument("--pos_file", required=True, help="Path to the positive sample TSV file.")
    parser.add_argument("--neg_file", required=True, help="Path to the negative sample TSV file.")
    parser.add_argument("--predictions_file", required=True, help="Path to the predictions TSV file.")
    parser.add_argument("--pos_output", required=True, help="Output file for the top positive sample keys.")
    parser.add_argument("--neg_output", required=True, help="Output file for the top negative sample keys.")
    parser.add_argument("--top_n", type=int, default=7500000, help="Number of top samples to select for each category.")

    args = parser.parse_args()

    main(args.pos_file, args.neg_file, args.predictions_file, args.pos_output, args.neg_output, args.top_n)